package com.ken.ability.gateway.loadbalance;

import com.alibaba.cloud.nacos.ribbon.NacosServer;
import com.ken.ability.gateway.application.GatewayGragProperties;
import com.netflix.client.config.IClientConfig;
import com.netflix.loadbalancer.BaseLoadBalancer;
import com.netflix.loadbalancer.ClientConfigEnabledRoundRobinRule;
import com.netflix.loadbalancer.Server;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.util.CollectionUtils;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 灰度发布的负载均衡器
 */
public class GrayLoadBalanceRule extends ClientConfigEnabledRoundRobinRule {


    //用于缓存不同版本轮询的下标值
    private AtomicInteger indexAtomic = new AtomicInteger(0);

    @Autowired
    private GatewayGragProperties gatewayGragProperties;

    @Override
    public Server choose(Object key) {

        //判断是否开启了灰度发布的功能，如果没开启则直接选择
        if (!gatewayGragProperties.isEnable()) {
            Server server = super.choose(key);
            return server;
        }

        //获得需要灰度的服务列表
        Map<String, GatewayGragProperties.GrayConfig> services = gatewayGragProperties.getServices();

        //获得负载均衡器
        BaseLoadBalancer loadBalancer = (BaseLoadBalancer) getLoadBalancer();

        //如果灰度的请求不包含当前服务，就直接沿用默认的负载均衡策略
        if (!services.containsKey(loadBalancer.getName())) {
            return super.choose(key);
        }

        //获得当前服务的灰度配置
        GatewayGragProperties.GrayConfig grayConfig = services.get(loadBalancer.getName());
        //获得需要灰度的版本
        String version = grayConfig.getVersion();

        //当前选中的Server对象
        Server server = null;
        //当前请求次数
        int count = 0;
        while (server == null && count++ < 10) {

            //通过负载均衡器获得对象
            List<Server> allServers = loadBalancer.getAllServers();
            System.out.println("获得所有负载均衡对象：" + allServers);
            List<Server> reachableServers = loadBalancer.getReachableServers();
            System.out.println("可用的负载均衡服务：" + reachableServers);

            int allCount = allServers.size();
            if (allCount == 0) {
                return null;
            }

            //符合当前版本的服务
            List<Server> runServer = new ArrayList<>();

            //根据灰度比例，是否开启灰度
            boolean isGray = Math.random() < grayConfig.getRate();

            //获得符合当前请求版本的服务
            for (Server aServer : allServers) {
                //转换成NacosServer
                NacosServer nacosServer = (NacosServer)aServer;
                //获得元数据信息
                Map<String, String> metadata = nacosServer.getMetadata();
                //判断当前服务的元数据信息中是否存在version
                if (metadata.containsKey("version")) {
                    //获得服务的版本号
                    String serverVersion = metadata.get("version");
                    if (isGray && Objects.equals(version, serverVersion) || !isGray && !Objects.equals(version, serverVersion)) {
                        //如果版本号匹配
                        runServer.add(aServer);
                    }
                }
            }

            //判断是否存在符合要求的服务器，没有则返回null
            if (CollectionUtils.isEmpty(runServer)) {
                return null;
            }

            //获得下一个服务下标
            int next = getNextServerIndex(runServer.size());
            server = runServer.get(next);

            if (server == null) {
                /* Transient. */
                Thread.yield();
                continue;
            }

            if (server.isAlive() && (server.isReadyToServe())) {
                return server;
            }

            server = null;
        }

        return server;
    }

    /**
     * 获得下一个请求下标
     * @return
     */
    private Integer getNextServerIndex(int size){
        for(;;) {
            int value = indexAtomic.get();
            int next = (value + 1) % size;
            if (indexAtomic.compareAndSet(value, next))
                return next;
        }
    }

    @Override
    public void initWithNiwsConfig(IClientConfig iClientConfig) {

    }
}
